# data_management/persistence.py
import json
import pickle
from typing import Dict, Any
from pathlib import Path
from datetime import datetime
import sqlite3

try:
    from utils.context_manager import ContextManager
    from models.model_strategy_manager import ModelStrategyManager
except ImportError:
    # 备用导入路径
    from context_manager import ContextManager
    from model_strategy_manager import ModelStrategyManager

class DataPersistence:
    """数据持久化管理"""
    
    def __init__(self, base_path: str = "data"):
        self.base_path = Path(base_path)
        self.base_path.mkdir(exist_ok=True)
        
        # 初始化数据库
        self.db_path = self.base_path / "onpharos.db"
        self._init_database()
    
    def _init_database(self):
        """初始化数据库"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 创建训练记录表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS training_records (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                cycle INTEGER NOT NULL,
                performance REAL NOT NULL,
                reward REAL NOT NULL,
                survival_rate REAL NOT NULL,
                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 创建因子状态表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS factor_states (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                cycle INTEGER NOT NULL,
                factor_name TEXT NOT NULL,
                level INTEGER NOT NULL,
                health REAL NOT NULL,
                attack REAL NOT NULL,
                defense REAL NOT NULL,
                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 创建事件记录表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS event_logs (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                cycle INTEGER NOT NULL,
                event_type TEXT NOT NULL,
                participants TEXT NOT NULL,
                outcome TEXT NOT NULL,
                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        conn.commit()
        conn.close()
    
    def save_training_cycle(self, cycle_data: Dict):
        """保存训练周期数据"""
        # JSON格式保存
        json_path = self.base_path / "training" / f"cycle_{cycle_data['cycle_number']}.json"
        json_path.parent.mkdir(exist_ok=True)
        
        with open(json_path, 'w', encoding='utf-8') as f:
            json.dump(cycle_data, f, indent=2, ensure_ascii=False)
        
        # 数据库保存
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            INSERT INTO training_records (cycle, performance, reward, survival_rate)
            VALUES (?, ?, ?, ?)
        ''', (
            cycle_data['cycle_number'],
            cycle_data.get('overall_performance', 0),
            cycle_data.get('average_reward', 0),
            cycle_data.get('average_survival_rate', 0)
        ))
        
        conn.commit()
        conn.close()
    
    def save_factor_states(self, cycle: int, factors: list):
        """保存因子状态"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        for factor in factors:
            cursor.execute('''
                INSERT INTO factor_states (cycle, factor_name, level, health, attack, defense)
                VALUES (?, ?, ?, ?, ?, ?)
            ''', (
                cycle,
                factor.name,
                factor.level,
                factor.attributes['health'],
                factor.attributes['attack'], 
                factor.attributes['defense']
            ))
        
        conn.commit()
        conn.close()
    
    def save_model_checkpoint(self, model, cycle: int, performance: float):
        """保存模型检查点"""
        checkpoint_path = self.base_path / "models" / f"model_cycle_{cycle}.pth"
        checkpoint_path.parent.mkdir(exist_ok=True)
        
        torch.save({
            'cycle': cycle,
            'performance': performance,
            'model_state': model.state_dict(),
            'timestamp': datetime.now().isoformat()
        }, checkpoint_path)
    
    def load_training_history(self, limit: int = 100) -> list:
        """加载训练历史"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT cycle, performance, reward, survival_rate, timestamp
            FROM training_records 
            ORDER BY cycle DESC 
            LIMIT ?
        ''', (limit,))
        
        history = []
        for row in cursor.fetchall():
            history.append({
                'cycle': row[0],
                'performance': row[1],
                'reward': row[2],
                'survival_rate': row[3],
                'timestamp': row[4]
            })
        
        conn.close()
        return history
    
    def load_factor_development(self, factor_name: str, cycles: int = 50) -> list:
        """加载因子发展数据"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT cycle, level, health, attack, defense
            FROM factor_states 
            WHERE factor_name = ?
            ORDER BY cycle DESC 
            LIMIT ?
        ''', (factor_name, cycles))
        
        development = []
        for row in cursor.fetchall():
            development.append({
                'cycle': row[0],
                'level': row[1],
                'health': row[2],
                'attack': row[3],
                'defense': row[4]
            })
        
        conn.close()
        return development
    
    def save_system_context(self, context_manager: ContextManager):
        """保存系统上下文"""
        # 上下文管理器有自己的持久化机制，这里可以实现额外的保存逻辑
        pass
    
    def save_model_strategies(self, strategy_manager: ModelStrategyManager, cycle: int):
        """保存模型策略状态"""
        # 保存策略配置和状态
        strategy_data = {
            'cycle': cycle,
            'strategies': {name: strategy.value for name, strategy in strategy_manager.model_strategies.items()},
            'timestamp': datetime.now().isoformat()
        }
        
        # 保存到数据库或文件
        strategy_path = self.base_path / "models" / f"strategy_cycle_{cycle}.json"
        strategy_path.parent.mkdir(exist_ok=True)
        
        with open(strategy_path, 'w', encoding='utf-8') as f:
            json.dump(strategy_data, f, indent=2, ensure_ascii=False)

    def cleanup_old_data(self, keep_cycles: int = 100):
        """清理旧数据"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 获取当前最大周期
        cursor.execute('SELECT MAX(cycle) FROM training_records')
        max_cycle = cursor.fetchone()[0] or 0
        
        if max_cycle > keep_cycles:
            delete_before = max_cycle - keep_cycles
            
            # 删除旧记录
            cursor.execute('DELETE FROM training_records WHERE cycle < ?', (delete_before,))
            cursor.execute('DELETE FROM factor_states WHERE cycle < ?', (delete_before,))
            cursor.execute('DELETE FROM event_logs WHERE cycle < ?', (delete_before,))
            
            conn.commit()
        
        conn.close()
